# 使用 PyTorch 2.5.1 的基础镜像
ARG BASE_IMAGE=pytorch/pytorch:2.5.1-cuda11.8-cudnn9-runtime
FROM ${BASE_IMAGE}

# 设置环境变量
ENV APP_ROOT=/opt/sam2
ENV PYTHONUNBUFFERED=1
ENV SAM2_BUILD_CUDA=0

# 安装系统依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
    ffmpeg \
    libavutil-dev \
    libavcodec-dev \
    libavformat-dev \
    libswscale-dev \
    pkg-config \
    build-essential \
    libffi-dev

# 升级 pip 和 setuptools
RUN pip install --upgrade pip setuptools --index-url https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn

# 修复 ffmpeg 链接问题
RUN if [ -f /opt/conda/bin/ffmpeg ]; then rm /opt/conda/bin/ffmpeg; fi && ln -s /bin/ffmpeg /opt/conda/bin/ffmpeg

# 创建应用目录
RUN mkdir -p ${APP_ROOT}

# 复制 sam2 代码到容器中
COPY ./models/sam2 ${APP_ROOT}/agents/models/sam2

# 复制 mq 代码到容器中
COPY ./mq ${APP_ROOT}/agents/mq

# 设置工作目录
WORKDIR ${APP_ROOT}/agents/models/sam2

# 安装 Python 依赖，使用清华源
RUN pip install -r requirements.txt --no-deps --no-build-isolation -v --index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 下载 SAM 2.1 检查点（如果需要）
RUN mkdir -p ${APP_ROOT}/agents/models/sam2/checkpoints
ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt ${APP_ROOT}/agents/models/sam2/checkpoints/sam2.1_hiera_tiny.pt
ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt ${APP_ROOT}/agents/models/sam2/checkpoints/sam2.1_hiera_small.pt
ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt ${APP_ROOT}/agents/models/sam2/checkpoints/sam2.1_hiera_base_plus.pt
ADD https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt ${APP_ROOT}/agents/models/sam2/checkpoints/sam2.1_hiera_large.pt

# 设置启动命令
CMD ["python", "run.py"]